use crate::{Error, prelude::*};
use core::alloc::{GlobalAlloc, Layout};
use core::mem::{needs_drop, size_of, MaybeUninit};
use core::ptr::{self, NonNull};
use core::slice;

/// # Safety
/// 这里分配的内存需要保证生命周期一定小于Self
/// 建议使用Boxed的构建函数，它保证了这一点
pub unsafe trait Allocator {
    /// # Safety
    /// 用户保证分配的内存生命周期一定小于self
    unsafe fn alloc_buf<F>(&self, layout: Layout, f: F) -> Result<NonNull<[u8]>, Error>
    where
        F: FnOnce(NonNull<[u8]>) -> Result<(), Error>;

    /// # Safety
    /// ptr一定时alloc_buf分配的内存地址
    unsafe fn deallocate(&self, ptr: NonNull<[u8]>, layout: Layout);

    /// # Safety
    /// 用户保证分配的内存生命周期一定小于self
    unsafe fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, Error> {
        self.alloc_buf(layout, |_| Ok(()))
    }

    /// # Safety
    /// ptr一定时alloc_buf分配的内存地址
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let slice = slice::from_raw_parts(ptr, layout.size());
        self.deallocate(NonNull::from(slice), layout)
    }
    /// # Safety
    /// 用户保证分配的内存生命周期一定小于self
    unsafe fn alloc_then<T, F>(&self, f: F) -> Result<NonNull<T>, Error>
    where
        F: FnOnce(&mut MaybeUninit<T>) -> Result<(), Error>,
    {
        let layout = Layout::new::<T>();
        self.alloc_buf(layout, |ptr| {
            let mut uninit = ptr.cast::<MaybeUninit<T>>();
            unsafe { f(uninit.as_mut()) }
        })
        .map(|ptr| ptr.cast::<T>())
    }

    /// # Safety
    /// 用户保证分配的内存生命周期一定小于self
    unsafe fn alloc_slice_then<T, F>(&self, len: usize, mut f: F) -> Result<NonNull<[T]>, Error>
    where
        F: FnMut(usize, &mut MaybeUninit<T>) -> Result<(), Error>,
    {
        let layout = Layout::array::<T>(len).map_err(|_| Error::new(EINVAL))?;
        self.alloc_buf(layout, |ptr| {
            let ptr = ptr.cast::<MaybeUninit<T>>();
            for offset in 0..len {
                let uninit = unsafe { &mut *ptr.as_ptr().add(offset) };
                f(offset, uninit).map_err(|e| {
                    if needs_drop::<T>() {
                        for n in (0..offset).rev() {
                            unsafe { (*ptr.as_ptr().add(n)).assume_init_drop() };
                        }
                    }
                    e
                })?;
            }
            Ok(())
        })
        .map(|ptr| {
            let ptr = ptr.cast::<T>().as_ptr();
            NonNull::from(unsafe { slice::from_raw_parts(ptr, len) })
        })
    }

    /// # Safety
    /// 用户保证分配的内存生命周期一定小于self
    unsafe fn init<T>(&self, val: T) -> Result<NonNull<T>, Error> {
        self.alloc_then::<T, _>(|obj| {
            obj.write(val);
            Ok(())
        })
    }

    /// # Safety
    /// 用户保证分配的内存生命周期一定小于self
    unsafe fn init_slice<T: Clone>(&self, len: usize, val: T) -> Result<NonNull<[T]>, Error> {
        self.alloc_slice_then::<T, _>(len, |_, obj| {
            obj.write(val.clone());
            Ok(())
        })
    }
    /// # Safety
    /// 用户保证分配的内存生命周期一定小于self
    unsafe fn zeroed<T>(&self) -> Result<NonNull<MaybeUninit<T>>, Error> {
        self.alloc_then::<T, _>(|obj| {
            unsafe { ptr::write_bytes(obj, 0, size_of::<T>()) };
            Ok(())
        })
        .map(|ptr| ptr.cast::<MaybeUninit<T>>())
    }

    /// # Safety
    /// 用户保证分配的内存生命周期一定小于self
    unsafe fn zeroed_slice<T>(&self, len: usize) -> Result<NonNull<[MaybeUninit<T>]>, Error> {
        self.alloc_slice_then::<T, _>(len, |_, _| Ok(()))
            .map(|boxed| {
                let ptr = boxed.as_ptr() as *mut u8;
                unsafe { ptr.write_bytes(0, size_of::<T>() * len) };
                boxed
            })
            .map(|ptr| {
                let ptr = ptr.cast::<MaybeUninit<T>>().as_ptr();
                NonNull::from(unsafe { slice::from_raw_parts(ptr, len) })
            })
    }

    /// # Safety
    /// 用户保证分配的内存生命周期一定小于self
    unsafe fn uninit<T>(&self) -> Result<NonNull<MaybeUninit<T>>, Error> {
        self.alloc_then::<T, _>(|_| Ok(()))
            .map(|ptr| ptr.cast::<MaybeUninit<T>>())
    }

    /// # Safety
    /// 用户保证分配的内存生命周期一定小于self
    unsafe fn uninit_slice<T>(&self, len: usize) -> Result<NonNull<[MaybeUninit<T>]>, Error> {
        self.alloc_slice_then::<T, _>(len, |_, _| Ok(()))
            .map(|ptr| {
                let ptr = ptr.cast::<MaybeUninit<T>>().as_ptr();
                NonNull::from(unsafe { slice::from_raw_parts(ptr, len) })
            })
    }

    /// # Safety
    /// ptr必须是allocate分配的内存, 基于类型获取的Layout和分配时保持一致
    unsafe fn release<T>(&self, ptr: NonNull<T>) {
        let layout = Layout::new::<T>();
        let ptr = unsafe { slice::from_raw_parts(ptr.cast::<u8>().as_ptr(), layout.size()) };
        self.deallocate(NonNull::from(ptr), layout);
    }

    /// # Safety
    /// ptr必须是allocate分配的内存, Layout和分配时保持一致
    unsafe fn release_with<T: ?Sized>(&self, ptr: NonNull<T>, layout: Layout) {
        let ptr = unsafe { slice::from_raw_parts(ptr.cast::<u8>().as_ptr(), layout.size()) };
        self.deallocate(NonNull::from(ptr), layout);
    }

    /// # Safety
    /// ptr必须是allocate分配的内存, 基于类型获取的Layout和分配时保持一致
    unsafe fn release_slice<T>(&self, ptr: NonNull<[T]>) {
        let layout = Layout::array::<T>(unsafe { ptr.as_ref() }.len()).unwrap();
        let ptr = unsafe { slice::from_raw_parts(ptr.cast::<u8>().as_ptr(), layout.size()) };
        self.deallocate(NonNull::from(ptr), layout);
    }
}

#[link(name = "c")]
extern "C" {
    fn aligned_alloc(alignment: usize, size: usize) -> *mut u8;
    fn free(ptr: *mut u8);
}

#[derive(Copy, Clone, PartialEq, PartialOrd, Ord, Eq)]
pub struct PoolAlloc;

unsafe impl Send for PoolAlloc {}
unsafe impl Sync for PoolAlloc {}

unsafe impl GlobalAlloc for PoolAlloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        aligned_alloc(layout.align(), layout.size())
    }
    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
        free(ptr);
    }
}

unsafe impl Allocator for PoolAlloc {
    unsafe fn alloc_buf<F>(&self, layout: Layout, f: F) -> Result<NonNull<[u8]>, Error>
    where
        F: FnOnce(NonNull<[u8]>) -> Result<(), Error>,
    {
        let ptr = GlobalAlloc::alloc(self, layout);
        if !ptr.is_null() {
            let ptr = NonNull::from(slice::from_raw_parts(ptr, layout.size()));
            f(ptr).map_err(|e| {
                GlobalAlloc::dealloc(self, ptr.cast::<u8>().as_ptr(), layout);
                e
            })?;
            Ok(ptr)
        } else {
            Err(Error::last_error())
        }
    }
    unsafe fn deallocate(&self, ptr: NonNull<[u8]>, layout: Layout) {
        GlobalAlloc::dealloc(self, ptr.cast::<u8>().as_ptr(), layout)
    }
}

#[derive(Copy, Clone, PartialEq, PartialOrd, Ord, Eq)]
pub struct NullAlloc;

unsafe impl GlobalAlloc for NullAlloc {
    unsafe fn alloc(&self, _layout: Layout) -> *mut u8 {
        ptr::null_mut()
    }
    unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {}
}

unsafe impl Allocator for NullAlloc {
    unsafe fn alloc_buf<F>(&self, _layout: Layout, _f: F) -> Result<NonNull<[u8]>, Error>
    where
        F: FnOnce(NonNull<[u8]>) -> Result<(), Error>,
    {
        Err(Error::default())
    }
    unsafe fn deallocate(&self, _ptr: NonNull<[u8]>, _layout: Layout) {}
}
